# 加载预训练模型
from transformers import BertModel
import torch

# 定义训练设备
DEVICE=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_dir="./model/google-bert/bert-base-chinese/models--google-bert--bert-base-chinese/snapshots/c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f"
pretrained=BertModel.from_pretrained(model_dir).to(DEVICE)

# print(pretrained)
'''
BertModel(
  (embeddings): BertEmbeddings(
    (word_embeddings): Embedding(21128, 768, padding_idx=0)
    (position_embeddings): Embedding(512, 768)
    (token_type_embeddings): Embedding(2, 768)
    (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
    (dropout): Dropout(p=0.1, inplace=False)
  )
  ...
  (pooler): BertPooler(
    (dense): Linear(in_features=768, out_features=768, bias=True) 
    (activation): Tanh()
  )
)
'''
# 定义下游任务模型（将主干网络所提取的特征进行分类）
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 上一步 最后输出全连接模型 Linear(in_features=768, out_features=768, bias=True)
        self.fc = torch.nn.Linear(768, 2)
    # 训练推理
    def forward(self,input_ids,attention_mask,token_type_ids):
        # # 冻结BERT不参与训练（仅前向传播） 上游任务不参与训练 不参与反向传播，只参与前向训练
       with torch.no_grad():
           out=pretrained(input_ids=input_ids,attention_mask=attention_mask,token_type_ids=token_type_ids)
       # 下游任务参与训练
       # 关键修改：取[CLS] token的特征（每个句子的整体表示）
       cls_representation = out.last_hidden_state[:, 0, :]  # 形状 [batch_size, 768]

       # 通过分类层
       output = self.fc(cls_representation)
       return output  # 去掉softmax，CrossEntropyLoss自带